Minimum falling path sum II

Time: O(MxN); Space: O(1); hard

Given a square grid of integers arr, a falling path with non-zero shifts is a choice of exactly one element from each row of arr, such that no two elements chosen in adjacent rows are in the same column.

Return the minimum sum of a falling path with non-zero shifts.

Example 1:

Input: arr = [[1,2,3],[4,5,6],[7,8,9]]

Output: 13

Explanation:

  • The possible falling paths are:

    • [1,5,9], [1,5,7], [1,6,7], [1,6,8],

    • [2,4,8], [2,4,9], [2,6,7], [2,6,8],

    • [3,4,8], [3,4,9], [3,5,7], [3,5,9]

  • The falling path with the smallest sum is [1,5,7], so the answer is 13.

Constraints:

  • 1 <= len(arr) = len(arr[i]) <= 200

  • -99 <= arr[i][j] <= 99

Hints:

  1. Use dynamic programming.

  2. Let dp[i][j] be the answer for the first i rows such that column j is chosen from row i.

  3. Use the concept of cumulative array to optimize the complexity of the solution.

1. Dynamic programming [O(MxN), O(1)]

[1]:

import heapq
class Solution1(object):
    """
    Time: O(M*N)
    Space: O(1)
    """
    def minFallingPathSum(self, arr):
        """
        :type arr: List[List[int]]
        :rtype: int
        """
        for i in range(1, len(arr)):
            smallest_two = heapq.nsmallest(2, arr[i-1])

            for j in range(len(arr[0])):
                arr[i][j] += smallest_two[1] if arr[i-1][j] == smallest_two[0] else smallest_two[0]

        return min(arr[-1])

[2]:

s = Solution1()

arr = [[1,2,3],[4,5,6],[7,8,9]]
assert s.minFallingPathSum(arr) == 13